import os
import nn
import utils
import caffe
import argparse
import mxnet as mx
import numpy.testing as npt

def convert_model(network, epoch, size, root):
    if not os.path.exists(root):
        os.makedirs(root)
    symbol = 'models/mtcnn/{0}/{0}.json'.format(network)
    params = 'models/mtcnn/{}/{:0>3}.params'.format(network, epoch)
    prototxt = '{}/{}.prototxt'.format(root, network)
    model = '{}/{}.caffemodel'.format(root, network)

    net = nn.get_net(network)
    net.load_parameters(params, ctx=mx.cpu())
    # convert model
    utils.convert_net(symbol, prototxt, size)
    utils.convert_model(net, prototxt, model)
    net2 = caffe.Net(prototxt, weights=model, phase=caffe.TEST)
    bn_maps = utils.merge_bn_proto(prototxt)
    net2 = utils.merge_bn_model(net2, prototxt, bn_maps)
    utils.convert_bgr(net2)
    utils.convert_mv(net2)
    net2.save(model)
    # check forward
    net2 = caffe.Net(prototxt, weights=model, phase=caffe.TEST)
    data = mx.nd.random.uniform(0, 255, shape=(1,3,size,size)).asnumpy()
    net2.blobs['Input1'].data[...] = data
    mx_data = (data[:, (2,1,0), :, :] - 127.5) * 0.0078125
    mx_out = net(mx.nd.array(mx_data))
    cf_out = net2.forward()
    for k, v in cf_out.items():
        if k.startswith('Softmax'):
            npt.assert_allclose(v, mx_out[0].asnumpy(), atol=1e-5)
        elif 'Convolution' in k:
            npt.assert_allclose(v, mx_out[1].asnumpy(), atol=1e-5)
        else:
            raise KeyError('Unknown key: {}'.format(k))

def parse_args():
    parser = argparse.ArgumentParser(description='Convert MXNet model to Caffe model')
    parser.add_argument('network', type=str)
    parser.add_argument('epoch', type=int)
    parser.add_argument('size', type=int)
    parser.add_argument('--root', type=str, default="models/_caffe")
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    convert_model(args.network, args.epoch, args.size, args.root)
    
